
import { useState, useEffect } from 'react';
import { Brain, Play, Square, Upload, Download, BarChart3, Zap } from 'lucide-react';
import { Card } from '@/components/ui/card';
import { Button } from '@/components/ui/button';
import { Progress } from '@/components/ui/progress';
import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from '@/components/ui/select';
import { Tabs, TabsContent, TabsList, TabsTrigger } from '@/components/ui/tabs';
import { Badge } from '@/components/ui/badge';

export default function ModelTraining() {
  const [isTraining, setIsTraining] = useState(false);
  const [trainingProgress, setTrainingProgress] = useState(0);
  const [currentEpoch, setCurrentEpoch] = useState(0);
  const [modelType, setModelType] = useState('cnn');
  const [trainingLoss, setTrainingLoss] = useState<number[]>([]);
  const [validationLoss, setValidationLoss] = useState<number[]>([]);

  // Simulate training progress
  useEffect(() => {
    if (isTraining) {
      const interval = setInterval(() => {
        setTrainingProgress(prev => {
          if (prev >= 100) {
            setIsTraining(false);
            return 100;
          }
          return prev + 2;
        });
        
        setCurrentEpoch(prev => prev + 1);
        
        // Simulate loss data
        if (Math.random() > 0.7) {
          setTrainingLoss(prev => [...prev.slice(-19), Math.random() * 0.5 + 0.1]);
          setValidationLoss(prev => [...prev.slice(-19), Math.random() * 0.6 + 0.15]);
        }
      }, 500);

      return () => clearInterval(interval);
    }
  }, [isTraining]);

  const startTraining = () => {
    setIsTraining(true);
    setTrainingProgress(0);
    setCurrentEpoch(0);
    setTrainingLoss([]);
    setValidationLoss([]);
  };

  const stopTraining = () => {
    setIsTraining(false);
  };

  return (
    <div className="space-y-6">
      {/* Header */}
      <div className="flex justify-between items-center">
        <div>
          <h2 className="text-2xl font-bold text-white">模型训练</h2>
          <p className="text-gray-400">深度学习模型训练与优化</p>
        </div>
        <div className="flex space-x-3">
          <Button variant="outline" size="sm">
            <Upload className="w-4 h-4 mr-2" />
            导入数据
          </Button>
          <Button
            onClick={isTraining ? stopTraining : startTraining}
            className={isTraining ? 'bg-red-500 hover:bg-red-600' : 'btn-primary'}
            disabled={isTraining && trainingProgress >= 100}
          >
            {isTraining ? (
              <>
                <Square className="w-4 h-4 mr-2" />
                停止训练
              </>
            ) : (
              <>
                <Play className="w-4 h-4 mr-2" />
                开始训练
              </>
            )}
          </Button>
        </div>
      </div>

      <div className="grid grid-cols-1 xl:grid-cols-3 gap-6">
        {/* Training Configuration */}
        <Card className="glass-card p-6">
          <h3 className="text-lg font-semibold mb-4 text-white flex items-center">
            <Brain className="w-5 h-5 mr-2" />
            训练配置
          </h3>
          
          <div className="space-y-4">
            <div>
              <label className="block text-sm font-medium text-gray-300 mb-2">模型类型</label>
              <Select value={modelType} onValueChange={setModelType}>
                <SelectTrigger className="input-advanced">
                  <SelectValue />
                </SelectTrigger>
                <SelectContent>
                  <SelectItem value="cnn">卷积神经网络 (CNN)</SelectItem>
                  <SelectItem value="rnn">循环神经网络 (RNN)</SelectItem>
                  <SelectItem value="transformer">Transformer</SelectItem>
                  <SelectItem value="resnet">ResNet</SelectItem>
                </SelectContent>
              </Select>
            </div>

            <div>
              <label className="block text-sm font-medium text-gray-300 mb-2">数据集</label>
              <Select defaultValue="custom">
                <SelectTrigger className="input-advanced">
                  <SelectValue />
                </SelectTrigger>
                <SelectContent>
                  <SelectItem value="custom">自定义数据集</SelectItem>
                  <SelectItem value="urbansound">UrbanSound8K</SelectItem>
                  <SelectItem value="esc50">ESC-50</SelectItem>
                </SelectContent>
              </Select>
            </div>

            <div>
              <label className="block text-sm font-medium text-gray-300 mb-2">训练参数</label>
              <div className="space-y-2 text-sm">
                <div className="flex justify-between">
                  <span className="text-gray-400">学习率</span>
                  <span className="text-green-400">0.001</span>
                </div>
                <div className="flex justify-between">
                  <span className="text-gray-400">批次大小</span>
                  <span className="text-green-400">32</span>
                </div>
                <div className="flex justify-between">
                  <span className="text-gray-400">总轮次</span>
                  <span className="text-green-400">100</span>
                </div>
              </div>
            </div>

            {isTraining && (
              <div className="pt-4 border-t border-gray-700">
                <div className="flex items-center justify-between mb-2">
                  <span className="text-sm text-gray-300">训练进度</span>
                  <Badge className="bg-blue-500/20 text-blue-400">
                    第 {currentEpoch} 轮
                  </Badge>
                </div>
                <Progress value={trainingProgress} className="mb-2" />
                <div className="text-center text-sm text-gray-400">
                  {trainingProgress.toFixed(1)}% 完成
                </div>
              </div>
            )}
          </div>
        </Card>

        {/* Training Visualization */}
        <div className="xl:col-span-2">
          <Card className="glass-card p-6 h-full">
            <h3 className="text-lg font-semibold mb-4 text-white flex items-center">
              <BarChart3 className="w-5 h-5 mr-2" />
              训练监控
            </h3>

            <Tabs defaultValue="loss" className="h-full">
              <TabsList className="grid w-full grid-cols-3">
                <TabsTrigger value="loss">损失曲线</TabsTrigger>
                <TabsTrigger value="accuracy">准确率</TabsTrigger>
                <TabsTrigger value="metrics">模型指标</TabsTrigger>
              </TabsList>
              
              <TabsContent value="loss" className="mt-4">
                <div className="h-64 bg-black/20 rounded-lg p-4">
                  <div className="flex justify-between mb-2">
                    <div className="text-sm">
                      <span className="text-blue-400">● 训练损失</span>
                      <span className="text-red-400 ml-4">● 验证损失</span>
                    </div>
                  </div>
                  <div className="h-full relative">
                    <svg width="100%" height="100%" viewBox="0 0 400 200">
                      {/* Training Loss */}
                      {trainingLoss.length > 1 && (
                        <path
                          d={`M ${trainingLoss.map((loss, i) => 
                            `${i * (380 / (trainingLoss.length - 1))},${180 - loss * 300}`
                          ).join(' L ')}`}
                          stroke="#3b82f6"
                          strokeWidth="2"
                          fill="none"
                        />
                      )}
                      {/* Validation Loss */}
                      {validationLoss.length > 1 && (
                        <path
                          d={`M ${validationLoss.map((loss, i) => 
                            `${i * (380 / (validationLoss.length - 1))},${180 - loss * 300}`
                          ).join(' L ')}`}
                          stroke="#ef4444"
                          strokeWidth="2"
                          fill="none"
                        />
                      )}
                    </svg>
                  </div>
                </div>
              </TabsContent>
              
              <TabsContent value="accuracy" className="mt-4">
                <div className="h-64 bg-black/20 rounded-lg p-4 flex items-center justify-center">
                  <div className="text-center">
                    <div className="text-4xl font-bold text-green-400 mb-2">
                      {isTraining ? `${(Math.random() * 20 + 75).toFixed(1)}%` : '92.3%'}
                    </div>
                    <div className="text-gray-400">当前准确率</div>
                  </div>
                </div>
              </TabsContent>
              
              <TabsContent value="metrics" className="mt-4">
                <div className="grid grid-cols-2 gap-4">
                  <div className="p-4 bg-black/20 rounded-lg">
                    <div className="text-purple-400 font-semibold mb-2">精确率</div>
                    <div className="text-2xl text-white">89.7%</div>
                  </div>
                  <div className="p-4 bg-black/20 rounded-lg">
                    <div className="text-pink-400 font-semibold mb-2">召回率</div>
                    <div className="text-2xl text-white">91.2%</div>
                  </div>
                  <div className="p-4 bg-black/20 rounded-lg">
                    <div className="text-yellow-400 font-semibold mb-2">F1 分数</div>
                    <div className="text-2xl text-white">90.4%</div>
                  </div>
                  <div className="p-4 bg-black/20 rounded-lg">
                    <div className="text-cyan-400 font-semibold mb-2">AUC</div>
                    <div className="text-2xl text-white">0.94</div>
                  </div>
                </div>
              </TabsContent>
            </Tabs>
          </Card>
        </div>
      </div>

      {/* Model Management */}
      <div className="grid grid-cols-1 md:grid-cols-2 gap-6">
        <Card className="glass-card p-6">
          <h3 className="text-lg font-semibold mb-4 text-white">模型管理</h3>
          <div className="space-y-3">
            {[
              { name: 'CNN_Model_v2.1', accuracy: '92.3%', size: '45MB', status: '已训练' },
              { name: 'RNN_Model_v1.8', accuracy: '88.7%', size: '32MB', status: '训练中' },
              { name: 'Transformer_v1.2', accuracy: '94.1%', size: '128MB', status: '已部署' },
            ].map((model, index) => (
              <div key={index} className="flex items-center justify-between p-3 bg-black/20 rounded-lg">
                <div>
                  <div className="font-medium text-white">{model.name}</div>
                  <div className="text-sm text-gray-400">
                    准确率: {model.accuracy} | 大小: {model.size}
                  </div>
                </div>
                <div className="flex items-center space-x-2">
                  <Badge className={
                    model.status === '已部署' ? 'bg-green-500/20 text-green-400' :
                    model.status === '训练中' ? 'bg-blue-500/20 text-blue-400' :
                    'bg-gray-500/20 text-gray-400'
                  }>
                    {model.status}
                  </Badge>
                  <Button size="sm" variant="outline">
                    <Download className="w-3 h-3" />
                  </Button>
                </div>
              </div>
            ))}
          </div>
        </Card>

        <Card className="glass-card p-6">
          <h3 className="text-lg font-semibold mb-4 text-white">系统资源</h3>
          <div className="space-y-4">
            <div>
              <div className="flex justify-between mb-2">
                <span className="text-gray-400">GPU 使用率</span>
                <span className="text-green-400">78%</span>
              </div>
              <Progress value={78} className="h-2" />
            </div>
            <div>
              <div className="flex justify-between mb-2">
                <span className="text-gray-400">内存使用</span>
                <span className="text-blue-400">12.5 / 16 GB</span>
              </div>
              <Progress value={78} className="h-2" />
            </div>
            <div>
              <div className="flex justify-between mb-2">
                <span className="text-gray-400">存储空间</span>
                <span className="text-yellow-400">145 / 500 GB</span>
              </div>
              <Progress value={29} className="h-2" />
            </div>
            <div className="pt-2 border-t border-gray-700">
              <div className="flex items-center justify-between">
                <span className="text-gray-400">预计完成时间</span>
                <span className="text-white flex items-center">
                  <Zap className="w-4 h-4 mr-1 text-yellow-400" />
                  {isTraining ? '25分钟' : '--'}
                </span>
              </div>
            </div>
          </div>
        </Card>
      </div>
    </div>
  );
}
